from sklearn.metrics import roc_auc_score
from tqdm import tqdm
from torch import nn
from easySample import easySample
from dataPre import dataPre
from torch.utils.data import Dataset, DataLoader
from siamese_graphsage_data import siamese_graphsage_dataset
import torch.nn.functional as F
import torch
from siamese_graphsage_model import SAGE
import torch_geometric.transforms as T
from torch_geometric.datasets import Planetoid
import sys
sys.path.append(r"/home/cyw/projects/function_sim_project/basic_script")

# 计算欧几里得距离
# euclidean_distance = torch.norm(vector1 - vector2, p=2)
#   下午就是细看一下模型的细节，确定怎么构建没有问题，然后就只剩下调通了
#   之后可以看一下预处理的asm2vec模型，别的数据来不急了，用我们之前处理的函数汇编信息吧。


class siamese_graphsage_trainer():
    def __init__(self, model) -> None:
        self.model = model
        self.epoch = 100
        self.lr = 0.001
        self.margin = 1
        # self.batch_size=1
        self.batch_size = 64
        self.optimizer = torch.optim.Adam(model.parameters(), lr=self.lr)

    def test_auc(self, dataloader, values):
        """
            用于查看验证集的auc，方便保存模型参数
        """
        ans = self.get_model_score(dataloader)
        auc_score = roc_auc_score(values, ans)
        return auc_score

    def train(self, dataloader, validdataloader, values):
        opcode = 0
        best_auc = -float("inf")
        for i in tqdm(range(self.epoch)):
            self.model.train()
            temploss = 0
            #   这个query想要处理成1，0的形式
            for j, (query, target1, target2) in enumerate(dataloader):
                opcode += 1
                self.optimizer.zero_grad()
                edge_index_x, att_x = target1
                edge_index_y, att_y = target2

                edge_index_x = edge_index_x.cuda()
                edge_index_y = edge_index_y.cuda()
                att_x = att_x.cuda()
                att_y = att_y.cuda()
                query = query.cuda()

                out = self.model(att_x, edge_index_x)
                out1 = self.model(att_y, edge_index_y)

                #   这个不支持batch
                #   Euclidean_distance=torch.norm(out - out1, p=2)
                euclidean_distance = nn.functional.pairwise_distance(out, out1)

                #   这个应该要反一下，我的相似值是1，不相似是0
                #   反了之后为啥auc还是小于0.5的
                # loss = torch.mean((1 - query) * 0.5 * torch.pow(euclidean_distance, 2) +
                #                 query * 0.5 * torch.pow(torch.clamp(self.margin - euclidean_distance, min=0.0), 2))
                loss = torch.mean(query * 0.5 * torch.pow(euclidean_distance, 2) +
                                  (1 - query) * 0.5 * torch.pow(torch.clamp(self.margin - euclidean_distance, min=0.0), 2))

                temploss += loss
                loss.backward()
                self.optimizer.step()
            temp_auc = self.test_auc(validdataloader, values)
            if temp_auc > best_auc:
                best_auc = temp_auc
                torch.save(
                    self.model, r'/home/cyw/projects/function_sim_project/all_data/models/siamese_graphsage_model.pth')
                print("best_auc:{}  保存模型数据".format(best_auc))
            print("epoch:{} loss = {}  valid auc = {}".format(
                i, temploss, temp_auc))
        print("siamese graphsage 训练完成!!")

    def get_model_score(self, dataloader):
        """
            这个模型的相似性是使用距离来衡量的，越相似的距离越近
            取值范围为【0，正无穷】
        """
        self.model.eval()
        with torch.no_grad():
            res = []
            for j, (query, target1, target2) in enumerate(dataloader):
                self.optimizer.zero_grad()
                edge_index_x, att_x = target1
                edge_index_y, att_y = target2

                edge_index_x = edge_index_x.cuda()
                edge_index_y = edge_index_y.cuda()
                att_x = att_x.cuda()
                att_y = att_y.cuda()
                query = query.cuda()

                out = self.model(att_x, edge_index_x)
                out1 = self.model(att_y, edge_index_y)

                euclidean_distance = nn.functional.pairwise_distance(out, out1)
                tempres = torch.clamp(1-euclidean_distance, min=0)
                res.append(tempres)
            ans = torch.cat((res))
        return ans.tolist()

    def get_model_dataloader(self, datas, name, is_shuffle=True):
        dataset = siamese_graphsage_dataset(datas, name)
        dataloader = DataLoader(dataset, batch_size=self.batch_size, shuffle=is_shuffle, num_workers=16,
                                collate_fn=dataset.adjust_samples_to_same_dimension)
        return dataloader


# nohup python /home/cyw/projects/function_sim_project/siamese_Graphsage/siamese_graphsage_train.py
if __name__ == "__main__":
    datapre = dataPre()
    data_pair_infs = datapre.get_pair_infs()
    values = []
    model = SAGE()
    s_sage = siamese_graphsage_trainer(model)
    dataloader = s_sage.get_model_dataloader(data_pair_infs["train"], "train")
    validdataloader = s_sage.get_model_dataloader(
        data_pair_infs["valid"], "valid", is_shuffle=False)
    for i in range(len(data_pair_infs["valid"])):
        values.append(data_pair_infs["valid"][i][2])
    s_sage.train(dataloader, validdataloader, values)

    # 简单测试能否跑通
    # 能正常训练
    # sample="387e8b94fb412139a600bd1c20c3b8b6"
    # testPair=[(sample,sample,1),
    #           ("1cf593553f5b14ec993125f146d35d72",sample,0),
    #           ("1cf593553f5b14ec993125f146d35d72","1cf593553f5b14ec993125f146d35d72",1)]
    # values=[]
    # model = SAGE()
    # s_sage=siamese_graphsage_trainer(model)
    # dataloader=s_sage.get_model_dataloader(testPair,"else")
    # for i in range(len(testPair)):
    #     values.append(testPair[i][2])
    # s_sage.train(dataloader,dataloader,values)
